import multiprocessing
import threading
from typing import Any, Callable, Iterable, Optional, Union

import matplotlib
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.widgets import Button

from .repeating_timer import RepeatingTimer
from .utils import convert_to_iterator

BarsPoint = Union[float, Iterable[float], Iterable[Iterable[float]]]
BarsBufferData = Optional[tuple[np.ndarray, float, float]]
BarsViewData = Optional[tuple[np.ndarray, float, float]]


def plot_summary(data: BarsViewData, title: Optional[str]):
    ys_heights, ys_bottom, ys_top = data

    plt.figure(figsize=(8, 4))
    axes: plt.Axes = plt.gca()
    axes.grid()

    if title is not None:
        axes.set_title(title)

    # draw bars
    for y_heights in ys_heights:
        axes.bar(np.arange(np.shape(y_heights)[0]), y_heights, 1,
            edgecolor="black", linewidth=1, zorder=2)

    # adjust y range
    if ys_bottom == 0:
        axes.set_ylim(0, ys_top * 1.1)
    else:
        padding = (ys_top - ys_bottom) / 8
        axes.set_ylim(ys_bottom - padding, ys_top + padding)

    plt.show()


class ViewBarsBuffer:
    def __init__(self):
        self._lock = threading.RLock()
        self._condition = threading.Condition(self._lock)

        self._version = 0
        self._ys: Optional[np.ndarray] = None
        self._ys_min: float = 0
        self._ys_max: float = 0

    def append(self, ys: BarsPoint):
        ys = np.asarray(ys, dtype=np.float32)
        if ys.ndim == 0:
            ys = ys.reshape([1, 1])
        elif ys.ndim == 1:
            ys = ys.reshape([1, -1])
        elif ys.ndim == 2:
            pass
        else:
            raise ValueError(f"Expected ys.ndim = 0, 1 or 2, got ys.shape={ys.shape}")

        with self._lock:
            self._version += 1
            self._ys = ys
            self._ys_min = np.minimum(self._ys_min, np.min(ys))
            self._ys_max = np.maximum(self._ys_max, np.max(ys))
            self._condition.notify_all()

    def clear(self):
        with self._lock:
            self._version += 1
            self._ys = None
            self._ys_min = 0
            self._ys_max = 0
            self._condition.notify_all()

    def read(self) -> tuple[int, BarsBufferData]:
        with self._lock:
            if self._ys is None:
                return self._version, None
            return self._version, (self._ys, self._ys_min, self._ys_max)

    def wait_and_read(self, version: int) -> tuple[int, BarsBufferData]:
        with self._lock:
            while version == self._version:
                self._condition.wait()
            return self.read()


class ViewBarsView:
    def __init__(self,
        figure: plt.Figure, *,
        title: Optional[str] = None,
    ):
        self._figure = self._init_figure(figure)
        self._axes = self._init_axes(figure, title)
        self._play_button, self._clear_button = self._init_buttons(figure)
        self._play_button.on_clicked(lambda _: self._on_play_clicked())
        self._clear_button.on_clicked(lambda _: self._on_clear_clicked())

        self._ys_bars: Optional[list[Iterable[plt.Rectangle]]] = None

        self._paused = True
        self.request_pause: Optional[Callable[[], Any]] = None
        self.request_resume: Optional[Callable[[], Any]] = None
        self.request_clear: Optional[Callable[[], Any]] = None

    @staticmethod
    def _init_figure(figure: plt.Figure):
        figure.set_size_inches(8, 4)
        return figure

    @staticmethod
    def _init_axes(figure: plt.Figure, title: Optional[str] = None) -> plt.Axes:
        axes = figure.gca()
        axes.grid()

        if title is not None:
            axes.set_title(title)

        return axes

    @staticmethod
    def _init_buttons(figure: plt.Figure) -> tuple[plt.Button, plt.Button]:
        figure.subplots_adjust(0.1, 0.1, 0.95, 0.85)
        play_button = Button(figure.add_axes([0.01, 0.9, 0.09, 0.09]), "resume")
        clear_button = Button(figure.add_axes([0.11, 0.9, 0.09, 0.09]), "clear")
        return play_button, clear_button

    @staticmethod
    def _init_bars(axes: plt.Axes, ys_n: int, bars_n: int) -> list[tuple[plt.Rectangle, ...]]:
        return [axes.bar(
            np.zeros([bars_n]), np.zeros([bars_n]),
            edgecolor="black", linewidth=1,
            zorder=2 + (ys_n - yi) / ys_n
        ) for yi in range(ys_n)]

    def update(self, data: BarsViewData):
        if data is None:  # clear
            self._axes.cla()
            self._axes.grid()
            self._ys_bars = None
            self._figure.canvas.draw_idle()
            return

        ys, ys_min, ys_max = data
        if self._ys_bars is None:
            ys_n, bars_n = np.shape(ys)
            if ys_n == 0 or bars_n == 0:
                return
            self._ys_bars = self._init_bars(self._axes, ys_n, bars_n)
            self._axes.set_xlim(-0.5, bars_n - 0.5)

        # update bars
        for y_bars, y_heights in zip(self._ys_bars, ys):
            for i, (bar, height) in enumerate(zip(y_bars, y_heights)):
                bar.set_x(i - 0.5)
                bar.set_width(1)
                bar.set_height(height)

        # adjust y range
        if ys_min == 0:
            self._axes.set_ylim(0, ys_max * 1.1)
        else:
            padding = (ys_max - ys_min) / 8
            self._axes.set_ylim(ys_min - padding, ys_max + padding)

        # redraw
        self._figure.canvas.draw_idle()

    def pause(self):
        if self._paused:
            return
        self._paused = True
        self._play_button.label.set_text("resume")

    def resume(self):
        if not self._paused:
            return
        self._paused = False
        self._play_button.label.set_text("pause")

    def _on_play_clicked(self):
        if self._paused:
            if self.request_resume:
                self.request_resume()
        else:
            if self.request_pause:
                self.request_pause()

    def _on_clear_clicked(self):
        if self.request_clear:
            self.request_clear()


class ViewBarsPopup:
    def __init__(self,
        producer: Union[Iterable[BarsPoint], Callable[[], BarsPoint]], *,
        produce_interval: float = 1 / 128,
        frame_interval: float = 1 / 32,
        title: Optional[str] = None,
        start: bool = True,
    ):
        """ View a data point as bars in a popup window

        :param producer: an iterator (generator) or callable that produces data point with shape (ys_n, bars_n)
        :param produce_interval: interval between calls of `producer`
        :param frame_interval: interval between gui updates
        :param title: the title of the plot
        :param start: whether to start immediately
        """
        self._iterator = convert_to_iterator(producer)
        self._data_connection, sub_data_connection = multiprocessing.Pipe()
        self._event_connection, sub_event_connection = multiprocessing.Pipe()
        self._subprocess = multiprocessing.Process(
            target=self.subprocess_func,
            args=(sub_data_connection, sub_event_connection,
                  frame_interval, title))
        self._subprocess.start()
        self._buffer = ViewBarsBuffer()

        def tick_func():
            point = next(self._iterator)
            try:
                self._data_connection.send(point)
            except (OSError, BrokenPipeError):
                self._timer.stop()
                return
            self._buffer.append(point)

        self._timer = RepeatingTimer(produce_interval, tick_func)

        def event_thread_func():
            while True:
                try:
                    cmd = self._event_connection.recv()
                except EOFError:
                    break
                if cmd == "pause":
                    self.pause()
                elif cmd == "resume":
                    self.resume()
                elif cmd == "clear":
                    self.clear()

        self._event_thread = threading.Thread(target=event_thread_func)
        self._event_thread.start()

        self._title = title

        if start:
            self.resume()

    @staticmethod
    def subprocess_func(
        data_connection,
        event_connection,
        frame_interval: float,
        title: Optional[str],
    ):
        matplotlib.use('QtAgg')

        buffer = ViewBarsBuffer()

        def data_thread_func():
            while True:
                point = data_connection.recv()
                buffer.append(point)

        data_thread = threading.Thread(target=data_thread_func, daemon=True)
        data_thread.start()

        figure = plt.figure()
        view = ViewBarsView(figure, title=title)
        view_version = -1

        def tick_func():
            nonlocal view_version
            view_version, view_data = buffer.wait_and_read(view_version)
            view.update(view_data)

        timer = RepeatingTimer(frame_interval, tick_func)
        timer.start()

        view.request_pause = lambda: event_connection.send("pause")
        view.request_resume = lambda: event_connection.send("resume")
        view.request_clear = lambda: event_connection.send("clear")

        def event_thread_func():
            while True:
                cmd = event_connection.recv()
                if cmd == "pause":
                    view.pause()
                elif cmd == "resume":
                    view.resume()
                elif cmd == "clear":
                    buffer.clear()

        event_thread = threading.Thread(target=event_thread_func, daemon=True)
        event_thread.start()

        plt.show()

    def resume(self):
        self._timer.start()
        self._event_connection.send("resume")

    def pause(self):
        self._timer.stop()
        self._event_connection.send("pause")

    def clear(self):
        self._buffer.clear()
        self._event_connection.send("clear")

    def wait(self):
        if self._subprocess.is_alive():
            self._subprocess.join()
        self._data_connection.close()
        self._event_connection.close()
        self._timer.join()
        self._event_thread.join()

    def close(self):
        if self._subprocess.is_alive():
            self._subprocess.terminate()
            self._subprocess.join()
        self._data_connection.close()
        self._event_connection.close()
        self._timer.join()
        self._event_thread.join()

    def plot_summary(self):
        _, view_data = self._buffer.read()
        plot_summary(view_data, title=self._title)

    def wait_and_plot_summary(self):
        self.wait()
        self.plot_summary()

    def close_and_plot_summary(self):
        self.close()
        self.plot_summary()
